package unimelb.filesync.common.crypto;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.UnsupportedEncodingException;
import java.security.Key;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Arrays;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

import org.apache.commons.codec.binary.Base64;

public class CryptoUtils {
    
    private static final int HASH_ITERATION = 1024;

    private static final int VECTOR_SIZE = 16;
    
    protected static final String PREFIX_PUBLIC_KEY = ".pub";

    protected static final String PREFIX_PRIVATE_KEY = ".pri";

    private static final String AES_CIPHER_MODE = "AES/CBC/PKCS5PADDING";

    private static final String ALGORITHM_RSA = "RSA";

    private static final String ALGORITHM_AES = "AES";
    
    private final SecureRandom vectorGenerator;

    public CryptoUtils() {
        super();
        this.vectorGenerator = new SecureRandom();
    }
    
    public String sha1Hash(String message) {
        try {
            final MessageDigest sha1Digest = MessageDigest.getInstance("SHA-1");
            
            byte[] data = message.getBytes();
            for (int i = 1; i <= HASH_ITERATION; i++) {
                data = sha1Digest.digest(data);
                sha1Digest.reset();
            }
            
            return encode(data);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    public KeyPair createRsaKeyPair(int keySize) {
        try {
            final KeyPairGenerator kpg = KeyPairGenerator.getInstance(ALGORITHM_RSA);
            kpg.initialize(keySize);
            return kpg.genKeyPair();
        } catch (NoSuchAlgorithmException e) {
            throw new RuntimeException(e);
        }
    }

    public void saveRsaKeyPair(KeyPair keyPair, String fileName) {
        try {
            final KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM_RSA);
            final PKCS8EncodedKeySpec privateKeySpec = keyFactory.getKeySpec(keyPair.getPrivate(), PKCS8EncodedKeySpec.class);
            final X509EncodedKeySpec publicKeySpec = keyFactory.getKeySpec(keyPair.getPublic(), X509EncodedKeySpec.class);

            saveEncodedKey(fileName + PREFIX_PRIVATE_KEY, privateKeySpec.getEncoded());
            saveEncodedKey(fileName + PREFIX_PUBLIC_KEY, publicKeySpec.getEncoded());
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    public KeyPair readRsaKeyPair(String fileName) {
        try {
            final PKCS8EncodedKeySpec privateKeySpec = new PKCS8EncodedKeySpec(readEncodedKey(fileName + PREFIX_PRIVATE_KEY));
            final X509EncodedKeySpec publicKeySpec = new X509EncodedKeySpec(readEncodedKey(fileName + PREFIX_PUBLIC_KEY));

            final KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM_RSA);
            return new KeyPair(keyFactory.generatePublic(publicKeySpec), keyFactory.generatePrivate(privateKeySpec));
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    public PublicKey readRsaPublicKey(String fileName) {
        try {
            final X509EncodedKeySpec publicKeySpec = new X509EncodedKeySpec(readEncodedKey(fileName));
            return KeyFactory.getInstance(ALGORITHM_RSA).generatePublic(publicKeySpec);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    public void saveRsaPublicKey(String fileName, PublicKey key) {
        try {
            saveEncodedKey(fileName, key.getEncoded());
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    public String encodePublicKey(PublicKey key) {
        return encode(key.getEncoded());
    }

    public PublicKey decodePublicKey(String key) {
        try {
            final X509EncodedKeySpec publicKeySpec = new X509EncodedKeySpec(decode(key));
            return KeyFactory.getInstance(ALGORITHM_RSA).generatePublic(publicKeySpec);
        } catch (Exception e) {
            throw new RuntimeException(e);

        }
    }

    public byte[] encryptRsa(byte[] data, Key key) {
        try {
            final Cipher cipher = Cipher.getInstance(ALGORITHM_RSA);
            cipher.init(Cipher.ENCRYPT_MODE, key);
            return cipher.doFinal(data);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    public byte[] decryptRsa(byte[] data, Key key) {
        try {
            final Cipher cipher = Cipher.getInstance(ALGORITHM_RSA);
            cipher.init(Cipher.DECRYPT_MODE, key);
            return cipher.doFinal(data);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    public boolean isRsaPrivateKeyFileExist(String fileName) {
        return new File(fileName + PREFIX_PRIVATE_KEY).exists();
    }

    public boolean isRsaPublicKeyFileExist(String fileName) {
        return new File(fileName + PREFIX_PUBLIC_KEY).exists();
    }

    public boolean isRsaKeyFilesExist(String fileName) {
        return isRsaPrivateKeyFileExist(fileName) && isRsaPublicKeyFileExist(fileName);
    }

    public SecretKey createAesKey() {
        try {
            KeyGenerator generator = KeyGenerator.getInstance(ALGORITHM_AES);
            generator.init(128);
            return generator.generateKey();
        } catch (NoSuchAlgorithmException e) {
            throw new RuntimeException(e);
        }
    }

    public byte[] encryptAes(byte[] data, Key key) {
        try {
            final byte[] initializationVector = createInitializationVector();

            final Cipher cipher = Cipher.getInstance(AES_CIPHER_MODE);
            cipher.init(Cipher.ENCRYPT_MODE, key, new IvParameterSpec(initializationVector));
            return concateArray(initializationVector, cipher.doFinal(data));
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    public byte[] decryptAes(byte[] data, Key key) {
        try {
            final byte[] initVector = new byte[VECTOR_SIZE];
            System.arraycopy(data, 0, initVector, 0, VECTOR_SIZE);
            
            final Cipher cipher = Cipher.getInstance(AES_CIPHER_MODE);
            cipher.init(Cipher.DECRYPT_MODE, key, new IvParameterSpec(initVector));
            return cipher.doFinal(data, VECTOR_SIZE, data.length - VECTOR_SIZE);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    public String encodeSecretKey(SecretKey key) {
        return encode(key.getEncoded());
    }

    public SecretKey decodeSecretKey(String data) {
        try {
            return new SecretKeySpec(decode(data), ALGORITHM_AES);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    public String encode(byte[] data) {
        try {
            return new String(Base64.encodeBase64(data), "UTF-8");
        } catch (UnsupportedEncodingException e) {
            throw new RuntimeException(e);
        }
    }

    public byte[] decode(String data) {
        return Base64.decodeBase64(data);
    }

    private void saveEncodedKey(String fileName, byte[] encodedKey) throws IOException {
        FileOutputStream os = null;
        try {
            os = new FileOutputStream(fileName);
            os.write(encodedKey);
            os.flush();
        } finally {
            if (os != null) {
                os.close();
            }
        }
    }

    private byte[] readEncodedKey(String fileName) throws IOException {
        FileInputStream is = null;
        try {
            final File keyFile = new File(fileName);
            if (!keyFile.exists()) {
                throw new RuntimeException("Key file not found");
            }

            is = new FileInputStream(keyFile);
            final byte[] encodedKey = new byte[(int) keyFile.length()];
            is.read(encodedKey);
            return encodedKey;
        } finally {
            if (is != null) {
                is.close();
            }
        }
    }

    private byte[] createInitializationVector() {
        final byte[] initVector = new byte[VECTOR_SIZE];
        vectorGenerator.nextBytes(initVector);
        return initVector;
    }

    private byte[] concateArray(byte[] array1, byte[] array2) {
        final byte[] result = Arrays.copyOf(array1, array1.length + array2.length);
        System.arraycopy(array2, 0, result, array1.length, array2.length);
        return result;
    }
    
    public static void main(String[] args) {
        final CryptoUtils utils = new CryptoUtils();
        final String encryptPassword = utils.encode(utils.createAesKey().getEncoded());
        System.out.println(encryptPassword);
    }
}
